import os
os.environ["PROJ_LIB"] = r'C:\Users\ckpark\anaconda3\Library\share'

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import numpy.ma as ma
import scipy.stats as ss
import properscoring as ps
from scipy.stats import iqr
from scipy.stats import linregress
from pylab import plot, show, savefig, xlim, figure, \
                ylim, legend, boxplot, setp, axes
                
import random
import numpy as np
rng = np.random.default_rng()
from scipy.stats import bootstrap

def init_plotting():
    plt.rcParams['font.family'] = 'Times New Roman'
init_plotting()

def color_box(bp, color, alpha):

    # Define the elements to color. You can also add medians, fliers and means
    elements = ['boxes','caps','whiskers','fliers']

    # Iterate over each of the elements changing the color
    for elem in elements:
        [plt.setp(bp[elem][idx], color=color, alpha=alpha) for idx in range(len(bp[elem]))]
    return

nmme_list = ['CanCM4i','CanSIPS-IC3','GEM5-NEMO','GFDL-SPEAR','NASA-GEOSS2S','NCEP-CFSv2','RSMAS-CCSM4']

sc = []
for i in range(len(nmme_list)):
    read = np.loadtxt('../6.NMME/RESULT/NMME_'+nmme_list[i], ndmin=2, dtype='str')
    ens = np.array(read[:,0], dtype='int')  
    sc.append(len(ens))

nmme = np.zeros((len(nmme_list),np.max(sc),3))

for i in range(len(nmme_list)):
    read = np.loadtxt('../6.NMME/RESULT/NMME_'+nmme_list[i], ndmin=2, dtype='str')
    s_var = np.array(read[:,1:], dtype='float')  
    
    nmme[i,:sc[i],:] = s_var
    
snmme = np.zeros((len(nmme_list),np.max(sc),3))
for i in range(len(nmme_list)):
    for j in range(sc[i]):
        snmme[i,j,0] = nmme[i,j,0]
        snmme[i,j,1] = nmme[i,j,0] + nmme[i,j,1]
        snmme[i,j,2] = nmme[i,j,0] + nmme[i,j,1] + nmme[i,j,2]


adats = [132, 180, 240]
bdats = [198, 261, 312]

for k in range(3):
    cc = 0; a = 0; b = 0
    for t in range(len(nmme_list)):
        for x in range(sc[t]):
            if snmme[t,x,k] >= adats[k]:
                a += 1
            if snmme[t,x,k] >= bdats[k]:
                b += 1
            cc += 1
        
    print(k,'dro_week_ens : ',(a/cc)*100)
    print(k,'dro_recov_ens : ',(b/cc)*100)

        
nt = len(nmme_list)
fig = plt.figure(figsize=(12,8))
xlab = ['CanCM4i','CanSIPS'+'\n'+'-IC3','GEM5'+'\n'+'NEMO',\
        'GFDL'+'\n'+'SPEAR','NASA'+'\n'+'GEOSS2S','NCEP'+'\n'+'CFSv2','RSMAS'+'\n'+'CCSM4']
for xx in range(3):
    ax1 = fig.add_subplot(3,1,xx+1)

    xax = np.arange(0,nt)
    xticks = np.arange(0,nt)
    
    for t in range(nt):
       
        ameanlineprops = dict(linestyle='solid', linewidth=2.5, color='red')
        amedianprops = dict(linestyle='none', linewidth=2.5, color='black')
            
        amean = np.mean(snmme[t,:sc[t],xx])

        aflierprops = dict(color='black')
           
        violin_parts = ax1.violinplot(snmme[t,:sc[t],xx], positions =[t], showmeans=False,\
                       showmedians=False,showextrema=True)  
        
        ax1.scatter(t-0.008,amean,color='red',s=100,zorder=10)
        

        for partname in ('cbars','cmins','cmaxes'):
            vp = violin_parts[partname]
            vp.set_edgecolor('gray')
            vp.set_linewidth(2)
      
        for vp in violin_parts['bodies']:
            vp.set_facecolor('gray')
            #vp.set_edgecolor(rrred)
            #vp.set_linewidth(1)
            #vp.set_alpha(0.5)
            
    if xx == 0:       
        ax1.axhline(y=132, color='green', linestyle='solid')
        ax1.axhline(y=198, color='blue', linestyle='dashed')
    if xx == 1:       
        ax1.axhline(y=180, color='green', linestyle='solid')
        ax1.axhline(y=261, color='blue', linestyle='dashed')
    if xx == 2:       
        ax1.axhline(y=240, color='green', linestyle='solid')
        ax1.axhline(y=312, color='blue', linestyle='dashed')
 
    if xx == 2:
        ax1.set_xticklabels(xlab)
        ax1.set_xticks(xticks)
        ax1.set_xlabel('Model',fontsize=26)
    else:
        ax1.set_xticks(np.arange(-100,100,400))
    ax1.set_xlim(-0.5,6.5)
    

    ax1.set_yticks(np.arange(0,600,100))
    ax1.set_ylim(0,490)
    if xx == 1:
        ax1.set_ylabel('Accumulated'+'\n'+'precipitation (mm)',fontsize=26)

    #ax1.set_yticks(np.arange(-100,1400,1000))

    #ax1.set_ylim(0,500)
      
    ax1.tick_params(labelsize=20)
    #if pp == 0:
    #    ax1.legend(loc=4,fontsize=20)



plt.subplots_adjust(hspace=0.03,wspace=0.0)
plt.show()      
outfig = '../99.FIG/each_nmme_model_box.png'      
fig.savefig(outfig,bbox_inches='tight')
